import torch


def save_ckpt(ckpt, save_name):
    torch.save(ckpt, save_name)
    print('save changed ckpt to {}'.format(save_name))

"""
重置训练状态
"""
def reset_ckpt(ckpt):
    if 'scheduler' in ckpt:
        del ckpt['scheduler']
    if 'optimizer' in ckpt:
        del ckpt['optimizer']
    if 'iteration' in ckpt:
        ckpt['iteration'] = 0

"""
使用多分类器权重来初始化二分类器权重
"""
def initialize_bin_para(ckpt):
    # 获得原始分类器权重
    orig_weight = None

    param = "roi_heads.box_predictor.cls_score.weight"
    new_param = "roi_heads.box_predictor.bin_class_score.weight"

    if param in ckpt['model']:
        print("initialize with ", param)
        orig_weight = ckpt['model'][param]
        print("orig_weight: ", orig_weight.size())

    # 分类没有偏置
    # roi_heads.box_predictor.cls_score.weight   torch.Size([21, 1024])
    # roi_heads.box_predictor.bbox_pred.weight   torch.Size([80, 1024])
    # roi_heads.box_predictor.bbox_pred.bias   torch.Size([80])

    # 分组策略
    group_a = (0, 3, 5, 7, 12, 15, 16, 17, 18, 19)
    group_b = (1, 2, 4, 6, 8, 9, 10, 11, 13, 14)

    new_weight_a = orig_weight.index_select(0, torch.tensor(group_a).cuda())
    print("new_weight_a: ", new_weight_a.size())
    new_weight_b = orig_weight.index_select(0, torch.tensor(group_b).cuda())
    print("new_weight_b: ", new_weight_b.size())

    new_weight_a = torch.mean(new_weight_a, dim=0, keepdim=True)
    new_weight_b = torch.mean(new_weight_b, dim=0, keepdim=True)

    new_weight = torch.cat((new_weight_a, new_weight_b), dim=0)

    print("new_weight: ", new_weight.size())

    ckpt['model'][new_param] = new_weight


"""
将微调好的基类权重用未微调的基类权重替代，从而提高基类检测精度
效果一般
"""
def replace_finetune_with_base(ckpt_final, ckpt_base):
    for k in ckpt_final['model'].keys():
        if "cls_score" in k:
            print(k)

    print("----------------------\n")
    finetune_cls_weight = ckpt_final['model']['roi_heads.box_predictor.cls_score.weight']

    finetune_cls_bias = ckpt_final['model']['roi_heads.box_predictor.cls_score.bias']

    base_cls_weight = ckpt_base['model']['roi_heads.box_predictor.cls_score.weight']

    for k in ckpt_base['model'].keys():
        if "cls_score" in k:
            print(k)

    base_cls_bias = ckpt_base['model']['roi_heads.box_predictor.cls_score.bias']

    for i in [finetune_cls_bias, finetune_cls_weight, base_cls_bias, base_cls_weight]:
        print(i.size())
    prev_cls = base_cls_bias.size(0)
    prev_cls -= 1

    finetune_cls_weight[:prev_cls] = base_cls_weight[:prev_cls]
    finetune_cls_bias[:prev_cls] = base_cls_bias[:prev_cls]

    finetune_cls_weight[-1] = base_cls_weight[-1]
    finetune_cls_bias[-1] = base_cls_bias[-1]

    ckpt_final['model']['roi_heads.box_predictor.cls_score.weight'] = finetune_cls_weight
    ckpt_final['model']['roi_heads.box_predictor.cls_score.bias'] = finetune_cls_bias

    save_ckpt(ckpt_final, "replaced_best_model_final.pth")


if __name__ == "__main__":
    ckpt = torch.load("sp2_fc_1shot.pth")
    # initialize_bin_para(ckpt)
    #
    reset_ckpt(ckpt)
    save_ckpt(ckpt, "simple_reset_sp2_fc_1shot.pth")
